import png

class ImageWriter():
  def __init__(self, output_dir):
    self.output_dir = output_dir

  # Create 4 images (one test and three neighbours)
  def create_images(self, test_vector, train_vectors, neighbours):
    filename_t = self.output_dir + str(test_vector[1]) + '_%d.png'
    self.create_image(filename_t % 0, test_vector)
    idx = 1
    for neighbour in neighbours:
      self.create_image(filename_t % idx, train_vectors[neighbour])
      idx += 1

  # Write the image file.
  def create_image(self, filename, feature_vector):
    pixels = [int(x) for x in feature_vector[0]]
    pixels_32_32 = []
    indices = range(0, 1025, 32)
    for i in range(0, len(indices) - 1):
      pixels_32_32.append(pixels[indices[i]: indices[i+1]])

    png_file = open(filename, 'wb')
    png_writer = png.Writer(32, 32, greyscale=True, bitdepth=1)
    png_writer.write(png_file, pixels_32_32)
    png_file.close()
